use std::alloc::{self, Layout};
use std::ptr::{null_mut};
use std::mem;
use std::slice;

const ALLOC_TEST_MAGIC: u32 = 0x72ec82d2;
const MALLOC_PATTERN: u32 = 0xBAADF00D;
const FREE_PATTERN: u32 = 0xDEADBEEF;

#[repr(C)]
struct BlockHeader {
    magic_number: u32,
    bytes: usize,
}

static mut ALLOCATED_BYTES: usize = 0;
static mut ALLOCATION_LIMIT: i32 = -1;

fn alloc_test_get_header(ptr: *mut u8) -> *mut BlockHeader {
    unsafe {
        let result = (ptr as *mut BlockHeader).offset(-1);
        assert_eq!((*result).magic_number, ALLOC_TEST_MAGIC);
        result
    }
}

fn alloc_test_overwrite(ptr: *mut u8, length: usize, pattern: u32) {
    let byte_ptr = unsafe { slice::from_raw_parts_mut(ptr, length) };
    for (i, byte) in byte_ptr.iter_mut().enumerate() {
        let pattern_seq = i & 3;
        *byte = ((pattern >> (8 * pattern_seq)) & 0xff) as u8;
    }
}

fn alloc_test_malloc(bytes: usize) -> *mut u8 {
    unsafe {
        if ALLOCATION_LIMIT == 0 {
            return null_mut();
        }

        let layout = Layout::from_size_align_unchecked(bytes + mem::size_of::<BlockHeader>(), mem::align_of::<BlockHeader>());
        let ptr = alloc::alloc(layout);

        if ptr.is_null() {
            return null_mut();
        }

        let header = ptr.cast::<BlockHeader>();
        (*header).magic_number = ALLOC_TEST_MAGIC;
        (*header).bytes = bytes;

        let block_ptr = (header as *mut u8).offset(mem::size_of::<BlockHeader>() as isize);
        alloc_test_overwrite(block_ptr, bytes, MALLOC_PATTERN);

        ALLOCATED_BYTES += bytes;

        if ALLOCATION_LIMIT > 0 {
            ALLOCATION_LIMIT -= 1;
        }

        block_ptr
    }
}

fn alloc_test_free(ptr: *mut u8) {
    unsafe {
        if ptr.is_null() {
            return;
        }

        let header = alloc_test_get_header(ptr);
        let block_size = (*header).bytes;
        assert!(ALLOCATED_BYTES >= block_size);

        alloc_test_overwrite(ptr, block_size, FREE_PATTERN);
        (*header).magic_number = 0;

        let layout = Layout::from_size_align_unchecked(block_size + mem::size_of::<BlockHeader>(), mem::align_of::<BlockHeader>());
        alloc::dealloc(header as *mut u8, layout);

        ALLOCATED_BYTES -= block_size;
    }
}

fn alloc_test_realloc(ptr: *mut u8, bytes: usize) -> *mut u8 {
    unsafe {
        let new_ptr = alloc_test_malloc(bytes);
        if new_ptr.is_null() {
            return null_mut();
        }

        if !ptr.is_null() {
            let header = alloc_test_get_header(ptr);
            let bytes_to_copy = (*header).bytes.min(bytes);
            std::ptr::copy_nonoverlapping(ptr, new_ptr, bytes_to_copy);
            alloc_test_free(ptr);
        }

        new_ptr
    }
}

fn alloc_test_calloc(nmemb: usize, bytes: usize) -> *mut u8 {
    let total_bytes = nmemb * bytes;
    let result = alloc_test_malloc(total_bytes);
    if !result.is_null() {
        unsafe { std::ptr::write_bytes(result, 0, total_bytes) };
    }
    result
}

fn alloc_test_strdup(string: &str) -> *mut u8 {
    let len = string.len();
    let result = alloc_test_malloc(len + 1);
    if !result.is_null() {
        unsafe {
            std::ptr::copy_nonoverlapping(string.as_ptr(), result, len);
            *result.offset(len as isize) = 0;
        }
    }
    result
}

fn alloc_test_set_limit(alloc_count: i32) {
    unsafe {
        ALLOCATION_LIMIT = alloc_count;
    }
}

fn alloc_test_get_allocated() -> usize {
    unsafe { ALLOCATED_BYTES }
}

fn main() {
    // Example usage of the alloc_test functions
    let ptr = alloc_test_malloc(10);
    if !ptr.is_null() {
        println!("Allocated memory");
        alloc_test_free(ptr);
    }
}